#!/bin/bash

# 批量 SSH 脚本（优化版）：确保所有主机执行，支持并发，自动找 hosts.txt
# 需要安装sshpass
# 用法: ./batch_ssh.sh "df -h"

# =============== 配置 ===============
SSH_USER="root"
SSH_KEY_FILE="$HOME/.ssh/id_rsa"
MAX_CONCURRENT=5

# =============== 颜色 ===============
GREEN='\033[0;32m'; YELLOW='\033[1;33m'; RED='\033[0;31m'; NC='\033[0m'
log() { echo -e "${GREEN}[INFO] $1${NC}"; }
warn() { echo -e "${YELLOW}[WARN] $1${NC}"; }
error() { echo -e "${RED}[ERROR] $1${NC}"; }

# =============== 生成密钥 ===============
setup_ssh_key() {
    if [ -f "$SSH_KEY_FILE" ]; then
        log "SSH 私钥已存在: $SSH_KEY_FILE"
    else
        log "生成 SSH 密钥..."
        mkdir -p ~/.ssh
        ssh-keygen -t rsa -b 2048 -N "" -f "$SSH_KEY_FILE" -q && log "密钥生成成功"
    fi
}

# =============== 推送公钥 ===============
setup_passwordless() {
    local host=$1
    local password=$2
    if [ -n "$password" ] && command -v sshpass > /dev/null; then
        echo -e "\n[$host] 推送公钥..."
        sshpass -p "$password" ssh-copy-id -i "$SSH_KEY_FILE.pub" -o StrictHostKeyChecking=no -o ConnectTimeout=10 "${SSH_USER}@$host" &> /tmp/sshcopy_$host.log
        if [ $? -eq 0 ]; then
            log "公钥已推送到 $host"
        else
            error "推送公钥失败: $host"
            grep -i "error\|fail" /tmp/sshcopy_$host.log | head -n 3
        fi
    elif [ -n "$password" ]; then
        warn "sshpass 未安装，跳过 $host 的公钥推送"
    fi
}

# =============== 执行命令 ===============
execute_on_host() {
    local host=$1
    local cmd=$2
    echo "========== 执行于 $host =========="
    ssh -o StrictHostKeyChecking=no -o ConnectTimeout=10 -o BatchMode=yes \
        "${SSH_USER}@${host}" "$cmd" 2>&1 | sed "s/^/[$host] /"
    local exit_code=$?
    echo -e "\n"
    return $exit_code
}

# =============== 主逻辑 ===============
main() {
    local COMMAND="$1"
    if [ -z "$COMMAND" ]; then
        echo "用法: $0 \"命令\""
        exit 1
    fi

    # 自动查找 hosts.txt
    if [ -f "/opt/hosts.txt" ]; then
        HOSTS_FILE="/opt/hosts.txt"
        log "使用 /opt/hosts.txt"
    elif [ -f "./hosts.txt" ]; then
        HOSTS_FILE="./hosts.txt"
        log "使用当前目录 hosts.txt"
    else
        error "找不到 hosts.txt（请放在 /opt/ 或当前目录）"
        exit 1
    fi

    # 检查 hosts.txt 内容
    local host_count=$(grep -v -E '^#|^$' "$HOSTS_FILE" | wc -l)
    log "发现 $host_count 台主机需要处理"

    setup_ssh_key

    read -p "是否推送公钥？[y/N]: " -n 1 -r; echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        read -s -p "请输入默认密码 (如果hosts.txt中未指定): " HOST_PASSWORD; echo
    fi

    # 创建任务数组
    local hosts=()
    local passwords=()

    # 读取主机列表
    while IFS= read -r line; do
        # 跳过空行和注释行
        [[ -z "$line" || "$line" =~ ^#.*$ ]] && continue

        # 提取主机和密码
        if [[ "$line" == *:* ]]; then
            host=$(echo "$line" | cut -d: -f1)
            password=$(echo "$line" | cut -d: -f2-)
        else
            host="$line"
            password="$HOST_PASSWORD"
        fi

        hosts+=("$host")
        passwords+=("$password")
    done < "$HOSTS_FILE"

    # 推送公钥（如果需要）
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        log "开始推送公钥..."
        for i in "${!hosts[@]}"; do
            setup_passwordless "${hosts[i]}" "${passwords[i]}"
        done
    fi

    # 执行命令（使用更简单的并发控制）
    log "开始执行命令: $COMMAND"
    local success_count=0
    local failed_hosts=()
    local total_hosts=${#hosts[@]}

    # 使用临时文件记录结果
    local result_file=$(mktemp)

    # 使用简单的并发控制
    local current_jobs=0

    for host in "${hosts[@]}"; do
        # 如果当前任务数达到最大并发数，等待一个任务完成
        while [ $current_jobs -ge $MAX_CONCURRENT ]; do
            # 等待任意后台任务完成
            wait -n
            current_jobs=$((current_jobs - 1))
        done

        # 启动后台任务
        (
            execute_on_host "$host" "$COMMAND"
            exit_code=$?
            echo "$host:$exit_code" >> "$result_file"
        ) &

        current_jobs=$((current_jobs + 1))
        log "[$host] 任务已启动 (当前并发数: $current_jobs)"
    done

    # 等待所有剩余任务完成
    wait

    # 收集结果
    while IFS=: read -r host status; do
        if [ "$status" -eq 0 ]; then
            ((success_count++))
        else
            failed_hosts+=("$host (退出码: $status)")
        fi
    done < "$result_file"

    rm -f "$result_file"

    # 显示汇总结果
    log "✅ 所有主机任务执行完成。"
    log "📊 执行结果: $success_count/$total_hosts 成功"

    if [ ${#failed_hosts[@]} -gt 0 ]; then
        error "❌ 失败的主机:"
        for host_info in "${failed_hosts[@]}"; do
            error "   - $host_info"
        done
    fi
}

main "$@"
